package tree

type trieNode struct {
	/*
		当前节点是不是单词的结尾
	*/
	isWord bool

	/*
		当前节点指向的指针的集合
	*/
	nextMap map[rune]*trieNode
}

type Trie struct {
	/*
		字典树的根节点
	*/
	root *trieNode

	/*
		字典树中单词的数量
	*/
	size int
}

/*
	获取字典树中单词的数量
*/
func (trie Trie) GetSize() int {
	return trie.size
}

/*
	O(len(word))
	向字典树中添加单词
*/
func (trie *Trie) Add(word string) {
	charArr := []rune(word)
	cur := trie.root
	for _, v := range charArr {
		if cur.nextMap == nil {
			cur.nextMap = make(map[rune]*trieNode)
		}
		_, contains := cur.nextMap[v]
		if !contains {
			cur.nextMap[v] = &trieNode{}
		}
		cur = cur.nextMap[v]
	}
	if !cur.isWord {
		// 如果单词没有出现过size++
		trie.size++
		// 最后一个标识为单词的结尾
		cur.isWord = true
	}
}

/*
	O(len(word))
	查询字典树中是否包含某个单词
*/
func (trie Trie) Contains(word string) bool {
	charArr := []rune(word)
	cur := trie.root
	for _, v := range charArr {
		if cur.nextMap == nil {
			return false
		}
		if _, ok := cur.nextMap[v]; !ok {
			return false
		}
		cur = cur.nextMap[v]
	}
	return cur.isWord
}

/*
	O(len(word))
	查询字典树中是否包含某个字符串开头的字符串
*/
func (trie Trie) IsPrefix(prefix string) bool {
	charArr := []rune(prefix)
	cur := trie.root
	for _, v := range charArr {
		if cur.nextMap == nil {
			return false
		}
		if _, ok := cur.nextMap[v]; !ok {
			return false
		}
		cur = cur.nextMap[v]
	}
	return true
}

/*
	满足前缀的单词过多的时候,无法保证较高的效率
	返回所有以某个前缀开头的单词
*/
func (trie Trie) PrefixWords(prefix string) []string {
	var res []string
	trie.prefixWords(trie.root, []rune(prefix), 0, &res, prefix)
	return res
}

func (trie Trie) prefixWords(node *trieNode, runes []rune, index int, res *[]string, word string) {
	if node == nil {
		return
	}
	if index < len(runes) {
		if len(node.nextMap) == 0 {
			return
		}
		value, contains := node.nextMap[runes[index]]
		if !contains {
			return
		}
		// 满足前缀就继续匹配
		trie.prefixWords(value, runes, index+1, res, word)
	} else {
		// 前缀满足,并且是一个完整的单词
		if node.isWord {
			*res = append(*res, word)
		}
		if len(node.nextMap) == 0 {
			return
		}
		// 寻找所有路径满足前缀的单词
		for c, v := range node.nextMap {
			trie.prefixWords(v, runes, index, res, string(append([]rune(word), c)))
		}
	}
}

/*
	.过多的时候时效性无法保证
	模糊搜索字符串,可以使用.表示所有与字符.
*/
func (trie Trie) Match(word string) bool {
	runes := []rune(word)
	return trie.match(trie.root, runes, 0)
}

func (trie *Trie) match(node *trieNode, runes []rune, index int) bool {
	// 匹配结束,查看最后一个位置是不是一个单词
	if len(runes) == index {
		return node.isWord
	}
	c := runes[index]
	if '.' == c {
		// 遍历所有可能满足条件的路径
		for _, v := range node.nextMap {
			if trie.match(v, runes, index+1) {
				// 有一条路径满足
				return true
			}
		}
		// 所有路径都不满足
		return false
	} else {
		next, contains := node.nextMap[c]
		if contains {
			return trie.match(next, runes, index+1)
		} else {
			return false
		}
	}
}
